#include "testRaytracing.h"

#include <VecGeom/surfaces/cuda/BrepCudaManager.h>
#include <VecGeom/surfaces/Model.h>
#include <VecGeom/surfaces/Navigator.h>
#include <VecGeom/surfaces/BVHSurfNavigator.h>
#include <VecGeom/management/BVHManager.h>
#include <VecGeom/navigation/BVHNavigator.h>
#include <VecGeom/navigation/LoopNavigator.h>
#include <VecGeom/base/Stopwatch.h>

using namespace vecgeom;
// surface model is based on Real_t precision defined in testRaytracing.h
using BrepCudaManager = vgbrep::BrepCudaManager<Real_t>;
using SurfData        = vgbrep::SurfData<Real_t>;
using vecCore::math::Abs;
// generation of rays etc is based on vecgeom::Precision as these are used also in the solid model
using Vec3D  = vecgeom::Vector3D<vecgeom::Precision>;
using Vec3Dc = Precision[3];

//==================================================================================
__global__ void LocateSolids(Vector3D<Precision> const *points, NavigationState *in_states, const VPlacedVolume *world,
                             TestConfig const config)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
    Vector3D<Precision> const &pos = points[i];
    // Locate with solid-based model
    LoopNavigator::LocatePointIn(world, pos, in_states[i], true);
  }
}
//==================================================================================
__global__ void LocateSolidsBVH(Vector3D<Precision> const *points, NavigationState *in_states,
                                const VPlacedVolume *world, TestConfig const config)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
    Vector3D<Precision> const &pos = points[i];
    // Locate with solid-based model
    BVHNavigator::LocatePointIn(world, pos, in_states[i], true);
  }
}
//==================================================================================
__global__ void ValidateLocate(NavigationState const *in_states, NavigationState const *out_states, int *num_errors,
                               TestConfig const config)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
    if (out_states[i].GetNavIndex() != in_states[i].GetNavIndex()) atomicAdd(num_errors, 1);
  }
}
//==================================================================================
__global__ void ComputeSafetiesSurf(Vector3D<Precision> const *points, NavigationState *in_states, Precision *safeties,
                                    TestConfig const config)
{
  if (config.validate_results) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
      int exit_surf;
      safeties[i] = vgbrep::protonav::ComputeSafety<Precision, Real_t>(points[i], in_states[i], exit_surf);
    }
  } else {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
      int exit_surf;
      vgbrep::protonav::ComputeSafety<Precision, Real_t>(points[i], in_states[i], exit_surf);
    }
  }
}
//==================================================================================
__global__ void ComputeSafetiesSurfBVH(Vector3D<Precision> const *points, NavigationState *in_states,
                                       Precision *safeties, TestConfig const config)
{
  if (config.validate_results) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
      safeties[i] = vgbrep::protonav::BVHSurfNavigator<Real_t>::ComputeSafety(points[i], in_states[i]);
    }
  } else {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
      vgbrep::protonav::BVHSurfNavigator<Real_t>::ComputeSafety(points[i], in_states[i]);
    }
  }
}
//==================================================================================
__global__ void ComputeSafetiesSolid(Vector3D<Precision> const *points, NavigationState const *in_states,
                                     Precision *ref_safeties, TestConfig const config)
{
  if (config.validate_results) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
      ref_safeties[i] = LoopNavigator::ComputeSafety(points[i], in_states[i]);
    }
  } else {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
      LoopNavigator::ComputeSafety(points[i], in_states[i]);
    }
  }
}
//==================================================================================
__global__ void ComputeSafetiesSolidBVH(Vector3D<Precision> const *points, NavigationState const *in_states,
                                        Precision *ref_safeties, TestConfig const config)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
    BVHNavigator::ComputeSafety(points[i], in_states[i]);
  }
}
//==================================================================================
__global__ void ValidateSafety(Precision const *safeties, Precision const *refSafeties, int *num_better_safety,
                               int *num_worse_safety, TestConfig const config)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
    atomicAdd(num_better_safety, int(safeties[i] > refSafeties[i] + kToleranceBVH));
    atomicAdd(num_worse_safety, int(safeties[i] < refSafeties[i] - kToleranceBVH));
  }
}
//==================================================================================
template <typename Navigator>
__device__ void PropagateRaySolid(int i, Vector3D<Precision> const *points, Vector3D<Precision> const *dirs,
                                  NavigationState const *in_states, Precision *length_over_crossings, bool debug,
                                  TestConfig const config)
{
  constexpr double kPushDistance = 1000 * vecgeom::kToleranceDist<Precision>;
  if (debug) {
    printf("Device PropagateRaysSolid debug ray %d:\n", i);
    printf("   ");
    in_states[i].Print();
  }
  NavigationState start_state = in_states[i];
  NavigationState out_state;
  int num_cross   = 0;
  double dist_tot = 0;
  auto const &dir = dirs[i];
  auto pt         = points[i] + kTolerance * dir; // push the start and subsequent crossing points
  do {
    auto distance =
        Navigator::ComputeStepAndPropagatedState(pt, dir, kInfLength, start_state, out_state, kPushDistance);
    if (debug) {
      printf("     dist = %15.10f\n", distance);
      printf("   ");
      out_state.Print();
    }
    dist_tot += (num_cross + 1) * distance;
    pt += distance * dir;
    start_state = out_state;
    num_cross++;
    if (config.accept_zeros && distance < 1000 * vecgeom::kTolerance) num_cross--;
  } while (!out_state.IsOutside() && num_cross < config.max_cross);

  if (config.validate_results) length_over_crossings[i] = num_cross ? dist_tot / (num_cross + 1) : 0;
}
//==================================================================================
template <typename Navigator>
__global__ void PropagateRaysSolid(Vector3D<Precision> const *points, Vector3D<Precision> const *dirs,
                                   NavigationState const *in_states, Precision *length_over_crossings,
                                   TestConfig const config)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
    PropagateRaySolid<Navigator>(i, points, dirs, in_states, length_over_crossings, /*debug=*/false, config);
  }
}
//==================================================================================
__global__ void ValidateTraversal(Vector3D<Precision> const *points, Vector3D<Precision> const *dirs,
                                  NavigationState const *in_states, Precision *length_over_crossings,
                                  Precision *refLength_over_crossings, int *num_errors, const VPlacedVolume *world,
                                  const SurfData *surfdata, bool use_bvh, TestConfig const config)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
    bool error_dist =
        Abs(length_over_crossings[i] - refLength_over_crossings[i]) >
        vgbrep::RoundingError(static_cast<Precision>(refLength_over_crossings[i]), 1000 * kToleranceDist<Precision>);
    if (error_dist && config.debug && *num_errors == 0) {
      printf("Error on GPU with BVH: %i in ray %i: length_over_crossings[i] %f refLength_over_crossings[i] %f -point "
             "%.10f %.10f %.10f -direction %.10f %.10f %.10f \n",
             use_bvh, i, length_over_crossings[i], refLength_over_crossings[i], points[i][0], points[i][1],
             points[i][2], dirs[i][0], dirs[i][1], dirs[i][2]);
      // note that the printouts in the propagation would be a mess since there is no GPU synchronization in between.
      // printf("Replaying solid model...\n");
      // PropagateRaySolid<LoopNavigator>(i, points, dirs, in_states, refLength_over_crossings, debug, max_cross,
      // accept_zeros); printf("Replaying surface model with BVH %i", use_bvh); PropagateRaySurf(i, points, dirs,
      // in_states, length_over_crossings, world, surfdata, debug, max_cross,
      //                  use_bvh, accept_zeros);
    }
    atomicAdd(num_errors, int(error_dist));
  }
}
//==================================================================================
__global__ void ResetStates(int nrays, NavigationState *out_states)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < nrays; i += blockDim.x * gridDim.x)
    out_states[i].Clear();
}
//==================================================================================
__global__ void LocateSurf(Vector3D<Precision> const *points, NavigationState *out_states, TestConfig const config)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
    // Locate with surface-based model
    vgbrep::protonav::LocatePointIn<Precision, Real_t>(NavigationState::WorldId(), points[i], out_states[i], true);
  }
}
//==================================================================================
__global__ void LocateSurfBVH(Vector3D<Precision> const *points, NavigationState *out_states, TestConfig const config)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
    // Locate with surface-based model + BVH
    vgbrep::protonav::BVHSurfNavigator<Real_t>::LocatePointIn(NavigationState::WorldId(), points[i], out_states[i],
                                                              true);
  }
}
//==================================================================================
__device__ void PropagateRaySurf(int i, Vector3D<Precision> const *points, Vector3D<Precision> const *dirs,
                                 NavigationState const *in_states, Precision *length_over_crossings,
                                 const SurfData *surfdata, bool debug, TestConfig const &config)
{
  if (debug) {
    printf("PropagateRaysSurf debug ray %d:\n", i);
    printf("   ");
    in_states[i].Print();
  }
  NavigationState start_state = in_states[i];
  NavigationState out_state;
  int num_cross = 0;
  vgbrep::CrossedSurface crossed_surf;
  double dist_tot = 0;
  auto pt         = points[i];
  auto const &dir = dirs[i];
  do {
    crossed_surf.Set(0, 0, 0); // need to reset because the same inner tube surface can be crossed twice in a row

    Precision distance{0};
    // USING THE LOOPER
    distance = vgbrep::protonav::ComputeStepAndHit<Precision, Real_t>(pt, dir, start_state, out_state, crossed_surf);

    // exiting framed surface marked as overlapping, need to relocate
    if (surfdata->IsFSOverlapping(crossed_surf.hit_surf) && crossed_surf.hit_surf.GetFSindex() != -1) {
      vgbrep::protonav::ReLocatePointIn<Precision, Real_t>(start_state, pt + distance * dir, dir, out_state,
                                                           crossed_surf.exit_surf, distance);
    }
    if (crossed_surf.hit_surf.GetFSindex() == -1) {
      // Extruding overlap detected, relocating to correct starting state
      // Find true location for the crossing point
      NavigationState true_state;
      vgbrep::protonav::ReLocatePointIn<Precision, Real_t>(start_state, pt, dir, true_state, crossed_surf.exit_surf,
                                                           /*distance=*/Precision(0.));

      // Now replay to get correct distance
      distance = vgbrep::protonav::ComputeStepAndHit<Precision, Real_t>(pt, dir, true_state, out_state, crossed_surf);
      VECGEOM_ASSERT(distance != 0 && distance != vecgeom::InfinityLength<Precision>() &&
                     "Distance after relocation shouldn't be 0 or infinity");
    }
    if (debug) {
      printf("     dist = %.16f  surf = %d\n", distance, crossed_surf.hit_surf.common_id);
      printf("   ");
      out_state.Print();
    }
    dist_tot += (num_cross + 1) * distance;
    pt += distance * dir;
    start_state = out_state;
    num_cross++;
    if (config.accept_zeros && distance < 1000 * vecgeom::kTolerance) num_cross--;
  } while (!out_state.IsOutside() && num_cross < config.max_cross);

  if (config.validate_results) length_over_crossings[i] = num_cross ? dist_tot / (num_cross + 1) : 0;
}

//==================================================================================
__device__ void PropagateRaySurfBVH(int i, Vector3D<Precision> const *points, Vector3D<Precision> const *dirs,
                                    NavigationState const *in_states, Precision *length_over_crossings,
                                    const SurfData *surfdata, bool debug, TestConfig const &config)
{
  if (debug) {
    printf("PropagateRaysSurf debug ray %d:\n", i);
    printf("   ");
    in_states[i].Print();
  }
  NavigationState start_state = in_states[i];
  NavigationState out_state;
  int num_cross = 0;
  vgbrep::CrossedSurface crossed_surf;
  double dist_tot = 0;
  auto pt         = points[i];
  auto const &dir = dirs[i];
  do {
    crossed_surf.Set(0, 0, 0); // need to reset because the same inner tube surface can be crossed twice in a row

    Precision distance{0};
    // USING THE BVH
    distance =
        vgbrep::protonav::BVHSurfNavigator<Real_t>::ComputeStepAndHit(pt, dir, start_state, out_state, crossed_surf);

    // exiting framed surface marked as overlapping, need to relocate
    if (surfdata->IsFSOverlapping(crossed_surf.hit_surf) && crossed_surf.hit_surf.GetFSindex() != -1) {
      vgbrep::protonav::ReLocatePointIn<Precision, Real_t>(start_state, pt + distance * dir, dir, out_state,
                                                           crossed_surf.exit_surf, distance);
    }
    if (crossed_surf.hit_surf.GetFSindex() == -1) {
      // Extruding overlap detected, relocating to correct starting state
      // Find true location for the crossing point
      NavigationState true_state;
      vgbrep::protonav::ReLocatePointIn<Precision, Real_t>(start_state, pt, dir, true_state, crossed_surf.exit_surf,
                                                           /*distance=*/Precision(0.));

      // Now replay to get correct distance
      distance =
          vgbrep::protonav::BVHSurfNavigator<Real_t>::ComputeStepAndHit(pt, dir, true_state, out_state, crossed_surf);

      VECGEOM_ASSERT(distance != 0 && distance != vecgeom::InfinityLength<Precision>() &&
                     "Distance after relocation shouldn't be 0 or infinity");
    }
    if (debug) {
      printf("     dist = %.16f  surf = %d\n", distance, crossed_surf.hit_surf.common_id);
      printf("   ");
      out_state.Print();
    }
    dist_tot += (num_cross + 1) * distance;
    pt += distance * dir;
    start_state = out_state;
    num_cross++;
    if (config.accept_zeros && distance < 1000 * vecgeom::kTolerance) num_cross--;
  } while (!out_state.IsOutside() && num_cross < config.max_cross);

  if (config.validate_results) length_over_crossings[i] = num_cross ? dist_tot / (num_cross + 1) : 0;
}

__device__ void PropagateRaySurfBVHSingle(int i, Vector3D<Precision> *points, Vector3D<Precision> const *dirs,
                                          NavigationState *in_states, Precision *length_over_crossings, bool debug,
                                          TestConfig const &config)
{
  if (in_states[i].IsOutside()) return;
  if (debug) {
    printf("PropagateRaySurfBVHSingle debug ray %d:\n", i);
    printf("   ");
    in_states[i].Print();
  }
  NavigationState start_state = in_states[i];
  NavigationState out_state;
  vgbrep::CrossedSurface crossed_surf;
  // commented out, needed if validation should be added
  // double dist_tot = 0;
  // int num_cross = 0;
  auto &pt        = points[i];
  auto const &dir = dirs[i];

  Precision distance{0};
  // USING THE BVH
  distance =
      vgbrep::protonav::BVHSurfNavigator<Real_t>::ComputeStepAndHit(pt, dir, start_state, out_state, crossed_surf);

  if (debug) {
    printf("     dist = %.16f  surf = %d\n", distance, crossed_surf.hit_surf.common_id);
    printf("   ");
    out_state.Print();
  }
  // dist_tot += (num_cross + 1) * distance;
  points[i] += distance * dir;
  in_states[i] = out_state;
  // num_cross++;
  // if (accept_zeros && distance < 1000 * vecgeom::kTolerance) num_cross--;
  // if (validate_results) length_over_crossings[i] = num_cross ? dist_tot / (num_cross + 1) : 0;
}

__device__ void ComputeStepAndNextSurface(int i, Vector3D<Precision> *points, Vector3D<Precision> const *dirs,
                                          NavigationState *in_states, long *hitcandidate_index, Precision *distances,
                                          Precision *length_over_crossings, bool debug, TestConfig const &config)
{
  if (in_states[i].IsOutside()) return;
  if (debug) {
    printf("ComputeStepAndNextSurface debug ray %d:\n", i);
    printf("   ");
    in_states[i].Print();
  }
  NavigationState start_state = in_states[i];
  NavigationState out_state;
  // int num_cross = 0;
  vgbrep::CrossedSurface crossed_surf;
  // double dist_tot = 0;
  auto &pt        = points[i];
  auto const &dir = dirs[i];
  auto &hit_index = hitcandidate_index[i];
  auto &distance  = distances[i];

  distance =
      vgbrep::protonav::BVHSurfNavigator<Real_t>::ComputeStepAndNextSurface(pt, dir, start_state, out_state, hit_index);

  if (debug) {
    printf("     dist = %.16f\n", distance);
  }
}

__device__ void RelocateToNextVolume(int i, Vector3D<Precision> *points, Vector3D<Precision> const *dirs,
                                     NavigationState *in_states, long *hitcandidate_index, Precision *distances,
                                     Precision *length_over_crossings, bool debug, TestConfig const &config)
{
  if (in_states[i].IsOutside()) return;
  if (debug) {
    printf("RelocateToNextVolume debug ray %d:\n", i);
    printf("   ");
    in_states[i].Print();
  }
  NavigationState out_state = in_states[i];
  vgbrep::CrossedSurface crossed_surf;
  auto &pt        = points[i];
  auto const &dir = dirs[i];
  auto &hit_index = hitcandidate_index[i];
  auto &distance  = distances[i];

  vgbrep::protonav::BVHSurfNavigator<Real_t>::RelocateToNextVolume(pt, dir, distance, hit_index, out_state,
                                                                   crossed_surf);

  points[i] += distance * dir;
  in_states[i] = out_state;
}

//==================================================================================
__global__ void PropagateRaysSurf(Vector3D<Precision> const *points, Vector3D<Precision> const *dirs,
                                  NavigationState const *in_states, Precision *length_over_crossings,
                                  const SurfData *surfdata, bool debug, TestConfig const config)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
    PropagateRaySurf(i, points, dirs, in_states, length_over_crossings, surfdata, debug, config);
  }
}

//==================================================================================
__global__ void PropagateRaysSurfBVH(Vector3D<Precision> const *points, Vector3D<Precision> const *dirs,
                                     NavigationState const *in_states, Precision *length_over_crossings,
                                     const SurfData *surfdata, bool debug, TestConfig const config)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < config.nrays; i += blockDim.x * gridDim.x) {
    PropagateRaySurfBVH(i, points, dirs, in_states, length_over_crossings, surfdata, debug, config);
  }
}

__global__ void PropagateRaysSurfBVHSingle(int nrays, int *alive_indices, Vector3D<Precision> *points,
                                           Vector3D<Precision> const *dirs, NavigationState *in_states,
                                           Precision *length_over_crossings, bool debug, TestConfig const config)
{
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < nrays; i += blockDim.x * gridDim.x) {
    PropagateRaySurfBVHSingle(alive_indices[i], points, dirs, in_states, length_over_crossings, debug, config);
  }
}

__global__ void ComputeStepAndNextSurfaces(int nrays, int *alive_indices, Vector3D<Precision> *points,
                                           Vector3D<Precision> const *dirs, NavigationState *in_states,
                                           long *hitcandidate_index, Precision *distances,
                                           Precision *length_over_crossings, bool debug, TestConfig const config)
{

  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < nrays; i += blockDim.x * gridDim.x) {
    ComputeStepAndNextSurface(alive_indices[i], points, dirs, in_states, hitcandidate_index, distances,
                              length_over_crossings, debug, config);
  }
}

__global__ void RelocateToNextVolumes(int nrays, int *alive_indices, Vector3D<Precision> *points,
                                      Vector3D<Precision> const *dirs, NavigationState *in_states,
                                      long *hitcandidate_index, Precision *distances, Precision *length_over_crossings,
                                      bool debug, TestConfig const config)
{

  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < nrays; i += blockDim.x * gridDim.x) {
    RelocateToNextVolume(alive_indices[i], points, dirs, in_states, hitcandidate_index, distances,
                         length_over_crossings, debug, config);
  }
}

__global__ void checkInsideStatesKernel(int nrays, const vecgeom::NavStateTuple *in_states, bool *result)
{
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < nrays) {
    if (!in_states[i].IsOutside()) {
      *result = true;
    }
  }
}

__global__ void filterAliveRays(const NavigationState *in_states, const int *indices, int *alive_indices,
                                int *alive_count, int active_count)
{
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx >= active_count) return;

  int current_index = indices[idx];

  // Check if the ray is not outside
  if (!in_states[current_index].IsOutside()) {
    int out_idx            = atomicAdd(alive_count, 1);
    alive_indices[out_idx] = current_index; // Store the original index of the alive ray
  }
}

void PropagateRaysSurfBVH(int initThreads, Vector3D<Precision> *points, Vector3D<Precision> *dirs,
                          NavigationState *in_states, Precision *length_over_crossings, bool debug,
                          TestConfig const &config)
{
  // This function does traversal for nrays with the BVH in a single step mode, such that each step is done in a single
  // kernel, similar to what is done in AdePT in a real simulation
  int *indices;
  VECGEOM_DEVICE_API_CALL(Malloc(&indices, config.nrays * sizeof(int)));
  std::vector<int> host_indices(config.nrays);
  std::iota(host_indices.begin(), host_indices.end(), 0);

  VECGEOM_DEVICE_API_CALL(
      Memcpy(indices, host_indices.data(), config.nrays * sizeof(int), VECGEOM_DEVICE_API_SYMBOL(MemcpyHostToDevice)));

  int *alive_indices;
  VECGEOM_DEVICE_API_CALL(Malloc(&alive_indices, config.nrays * sizeof(int)));

  int *alive_count;
  VECGEOM_DEVICE_API_CALL(Malloc(&alive_count, sizeof(int)));
  VECGEOM_DEVICE_API_CALL(Memset(alive_count, 0, sizeof(int)));

  int host_alive_count = config.nrays;
  int num_cross        = 0;
  Stopwatch timer;
  double time_per_step = 0.;

  while (host_alive_count) {
    int initBlocks = (host_alive_count + initThreads - 1) / initThreads;

    timer.Start();
    PropagateRaysSurfBVHSingle<<<initBlocks, initThreads>>>(host_alive_count, indices, points, dirs, in_states,
                                                            length_over_crossings, false, config);
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());

    time_per_step = timer.Stop();
    if (config.verbosity > 3) {
      std::cout << "istep: " << num_cross << " number of alive tracks: " << host_alive_count
                << " time : " << time_per_step << std::endl;
    }

    VECGEOM_DEVICE_API_CALL(Memset(alive_count, 0, sizeof(int)));
    filterAliveRays<<<initBlocks, initThreads>>>(in_states, indices, alive_indices, alive_count, host_alive_count);
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
    VECGEOM_DEVICE_API_CALL(
        Memcpy(&host_alive_count, alive_count, sizeof(int), VECGEOM_DEVICE_API_SYMBOL(MemcpyDeviceToHost)));
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());

    // Swap pointers so that indices now points to the new alive_indices for the next iteration
    std::swap(indices, alive_indices);

    num_cross++;
  }

  VECGEOM_DEVICE_API_CALL(Free(indices));
  VECGEOM_DEVICE_API_CALL(Free(alive_indices));
  VECGEOM_DEVICE_API_CALL(Free(alive_count));
}

void PropagateRaysSurfBVHSplit(int initThreads, Vector3D<Precision> *points, Vector3D<Precision> *dirs,
                               NavigationState *in_states, Precision *length_over_crossings, bool debug,
                               TestConfig const &config)
{
  // This function does traversal for nrays with the BVH in a single step mode and splits the traversal kernel into the
  // approach , such that each step is done in a single kernel, similar to what is done in AdePT in a real simulation
  int *indices;
  // Initialize initial indices
  VECGEOM_DEVICE_API_CALL(Malloc(&indices, config.nrays * sizeof(int)));
  std::vector<int> host_indices(config.nrays);
  std::iota(host_indices.begin(), host_indices.end(), 0);
  VECGEOM_DEVICE_API_CALL(
      Memcpy(indices, host_indices.data(), config.nrays * sizeof(int), VECGEOM_DEVICE_API_SYMBOL(MemcpyHostToDevice)));
  int *alive_indices;
  VECGEOM_DEVICE_API_CALL(Malloc(&alive_indices, config.nrays * sizeof(int)));
  int *alive_count;
  VECGEOM_DEVICE_API_CALL(Malloc(&alive_count, sizeof(int)));
  VECGEOM_DEVICE_API_CALL(Memset(alive_count, 0, sizeof(int)));

  long *hitcandidate_index;
  VECGEOM_DEVICE_API_CALL(Malloc(&hitcandidate_index, config.nrays * sizeof(long)));
  VECGEOM_DEVICE_API_CALL(Memset(hitcandidate_index, 0, sizeof(long)));

  Precision *distances;
  VECGEOM_DEVICE_API_CALL(Malloc(&distances, config.nrays * sizeof(Precision)));
  VECGEOM_DEVICE_API_CALL(Memset(distances, 0, sizeof(Precision)));

  int host_alive_count = config.nrays;
  int num_cross        = 0;
  Stopwatch timer;
  double time_per_step = 0.;

  int Blocks;

  while (host_alive_count) {
    Blocks = (host_alive_count + initThreads - 1) / initThreads;
    timer.Start();

    ComputeStepAndNextSurfaces<<<Blocks, initThreads>>>(host_alive_count, indices, points, dirs, in_states,
                                                        hitcandidate_index, distances, length_over_crossings, debug,
                                                        config);
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
    RelocateToNextVolumes<<<Blocks, initThreads>>>(host_alive_count, indices, points, dirs, in_states,
                                                   hitcandidate_index, distances, length_over_crossings, debug, config);
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
    time_per_step = timer.Stop();
    if (config.verbosity > 3) {
      std::cout << "istep: " << num_cross << " number of alive tracks: " << host_alive_count
                << " time : " << time_per_step << std::endl;
    }

    VECGEOM_DEVICE_API_CALL(Memset(alive_count, 0, sizeof(int)));
    filterAliveRays<<<Blocks, initThreads>>>(in_states, indices, alive_indices, alive_count, host_alive_count);
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
    VECGEOM_DEVICE_API_CALL(
        Memcpy(&host_alive_count, alive_count, sizeof(int), VECGEOM_DEVICE_API_SYMBOL(MemcpyDeviceToHost)));
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());

    // Swap pointers so that indices now points to the new alive_indices for the next iteration
    std::swap(indices, alive_indices);

    num_cross++;
  }

  VECGEOM_DEVICE_API_CALL(Free(indices));
  VECGEOM_DEVICE_API_CALL(Free(alive_indices));
  VECGEOM_DEVICE_API_CALL(Free(alive_count));
  VECGEOM_DEVICE_API_CALL(Free(distances));
  VECGEOM_DEVICE_API_CALL(Free(hitcandidate_index));
}

//==================================================================================
int testRaytracingCUDA(Vec3Dc const *pointsc, Vec3Dc const *dirsc, const SurfData &surfdata, TestConfig &config)
{
  BrepCudaManager::Instance().TransferSurfData(surfdata);
  auto surfdata_D = BrepCudaManager::Instance().GetDevicePtr();

  auto pointsh = new Vec3D[config.nrays];
  auto dirsh   = new Vec3D[config.nrays];
  for (auto i = 0; i < config.nrays; ++i) {
    pointsh[i].Set(pointsc[i][0], pointsc[i][1], pointsc[i][2]);
    dirsh[i].Set(dirsc[i][0], dirsc[i][1], dirsc[i][2]);
  }

  // Allocate/copy data on device
  Vec3D *points;
  VECGEOM_DEVICE_API_CALL(Malloc(&points, config.nrays * sizeof(Vec3D)));
  VECGEOM_DEVICE_API_CALL(
      Memcpy(points, pointsh, config.nrays * sizeof(Vec3D), VECGEOM_DEVICE_API_SYMBOL(MemcpyHostToDevice)));
  Vec3D *dirs;
  VECGEOM_DEVICE_API_CALL(Malloc(&dirs, config.nrays * sizeof(Vec3D)));
  VECGEOM_DEVICE_API_CALL(
      Memcpy(dirs, dirsh, config.nrays * sizeof(Vec3D), VECGEOM_DEVICE_API_SYMBOL(MemcpyHostToDevice)));
  Vec3D *points_NC;
  VECGEOM_DEVICE_API_CALL(Malloc(&points_NC, config.nrays * sizeof(Vec3D)));
  VECGEOM_DEVICE_API_CALL(
      Memcpy(points_NC, pointsh, config.nrays * sizeof(Vec3D), VECGEOM_DEVICE_API_SYMBOL(MemcpyHostToDevice)));
  NavigationState *origStates;
  VECGEOM_DEVICE_API_CALL(Malloc(&origStates, config.nrays * sizeof(NavigationState)));
  NavigationState *outputStates;
  VECGEOM_DEVICE_API_CALL(Malloc(&outputStates, config.nrays * sizeof(NavigationState)));
  NavigationState *outputStatesBVH;
  VECGEOM_DEVICE_API_CALL(Malloc(&outputStatesBVH, config.nrays * sizeof(NavigationState)));
  Precision *refSafeties{nullptr};
  Precision *safeties{nullptr};
  Precision *bvhSafeties{nullptr};
  Precision *refLength_over_crossings{nullptr};
  Precision *length_over_crossings{nullptr};
  Precision *length_over_crossings_bvh{nullptr};

  if (config.only_surf) config.debug = false;
  if (config.validate_results) {
    VECGEOM_DEVICE_API_CALL(Malloc(&refSafeties, config.nrays * sizeof(Precision)));
    VECGEOM_DEVICE_API_CALL(Malloc(&safeties, config.nrays * sizeof(Precision)));
    VECGEOM_DEVICE_API_CALL(Malloc(&bvhSafeties, config.nrays * sizeof(Precision)));
    VECGEOM_DEVICE_API_CALL(Malloc(&refLength_over_crossings, config.nrays * sizeof(Precision)));
    VECGEOM_DEVICE_API_CALL(Malloc(&length_over_crossings, config.nrays * sizeof(Precision)));
    VECGEOM_DEVICE_API_CALL(Malloc(&length_over_crossings_bvh, config.nrays * sizeof(Precision)));
  }

  constexpr int initThreads = 32;
  int initBlocks            = (config.nrays + initThreads - 1) / initThreads;

  Stopwatch timer;
  double time_locate_solids{0}, time_locate_solids_bvh{0}, time_locate_surf{0}, time_locate_surf_bvh{0},
      time_safety_solids{0}, time_safety_solids_bvh{0}, time_safety_surf{0}, time_safety_surf_bvh{0},
      time_traverse_solids{0}, time_traverse_solids_bvh{0}, time_traverse_surf{0}, time_traverse_surf_bvh{0};
  int num_errors          = 0;
  int num_errors_bvh_loc  = 0;
  int num_better_safety   = 0;
  int num_worse_safety    = 0;
  int num_errors_dist     = 0;
  int num_errors_dist_bvh = 0;
  int *num_errors_d, *num_better_safety_d, *num_worse_safety_d, *num_errors_dist_d, *num_errors_dist_bvh_d;
  VECGEOM_DEVICE_API_CALL(Malloc(&num_errors_d, sizeof(int)));
  VECGEOM_DEVICE_API_CALL(Malloc(&num_better_safety_d, sizeof(int)));
  VECGEOM_DEVICE_API_CALL(Malloc(&num_worse_safety_d, sizeof(int)));
  VECGEOM_DEVICE_API_CALL(Malloc(&num_errors_dist_d, sizeof(int)));
  VECGEOM_DEVICE_API_CALL(Malloc(&num_errors_dist_bvh_d, sizeof(int)));

  // Locating the global points with solid model
  if (!config.only_surf) {
    auto world_dev = vecgeom::cxx::CudaManager::Instance().world_gpu();
    timer.Start();
    LocateSolids<<<initBlocks, initThreads>>>(points, origStates, world_dev, config);
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
    time_locate_solids = timer.Stop();

    // Locating the global points with solid model + BVH
    timer.Start();
    LocateSolidsBVH<<<initBlocks, initThreads>>>(points, outputStates, world_dev, config);
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
    time_locate_solids_bvh = timer.Stop();

    ResetStates<<<initBlocks, initThreads>>>(config.nrays, outputStates);
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
  }

  // Locating the global points with surface model
  timer.Start();
  if (config.only_surf)
    LocateSurf<<<initBlocks, initThreads>>>(points, origStates, config);
  else
    LocateSurf<<<initBlocks, initThreads>>>(points, outputStates, config);
  VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
  time_locate_surf = timer.Stop();

  // Locating the global points with surface model + BVH
  timer.Start();
  if (config.test_bvh) LocateSurfBVH<<<initBlocks, initThreads>>>(points, outputStatesBVH, config);
  VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
  time_locate_surf_bvh = timer.Stop();

  if (!config.only_surf) {
    // Corectness for locating points
    ValidateLocate<<<initBlocks, initThreads>>>(origStates, outputStates, num_errors_d, config);
    VECGEOM_DEVICE_API_CALL(
        Memcpy(&num_errors, num_errors_d, sizeof(int), VECGEOM_DEVICE_API_SYMBOL(MemcpyDeviceToHost)));
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());

    if (num_errors > 0) std::cout << "Device: Point locate errors: " << num_errors << "\n";

    // Validate BVH Locate
    if (config.test_bvh) {
      ValidateLocate<<<initBlocks, initThreads>>>(origStates, outputStatesBVH, num_errors_d, config);
      VECGEOM_DEVICE_API_CALL(
          Memcpy(&num_errors_bvh_loc, num_errors_d, sizeof(int), VECGEOM_DEVICE_API_SYMBOL(MemcpyDeviceToHost)));
      VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
      if (num_errors_bvh_loc > 0) std::cout << "Device: BVH point locate errors: " << num_errors_bvh_loc << "\n";
    }
  }

  if (!config.debug) {
    if (config.only_surf)
      std::cout << "Device: locate_surf: " << time_locate_surf;
    else
      std::cout << "Device: locate_solids: " << time_locate_solids << "  locate_solids_BVH: " << time_locate_solids_bvh
                << "  locate_surf: " << time_locate_surf;
    if (config.test_bvh) {
      std::cout << " locate_surf_BVH: " << time_locate_surf_bvh << "\n";
    } else {
      std::cout << "\n";
    }
  }

  if (!config.only_surf) {
    // Safety for solids model (reference)
    timer.Start();
    ComputeSafetiesSolid<<<initBlocks, initThreads>>>(points, origStates, refSafeties, config);
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
    time_safety_solids = timer.Stop();

    // Safety for solids model with BVH
    timer.Start();
    ComputeSafetiesSolidBVH<<<initBlocks, initThreads>>>(points, origStates, safeties, config);
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
    time_safety_solids_bvh = timer.Stop();
  }

  // Safety for surface model
  timer.Start();
  ComputeSafetiesSurf<<<initBlocks, initThreads>>>(points, origStates, safeties, config);
  VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
  time_safety_surf = timer.Stop();

  timer.Start();
  if (config.test_bvh) ComputeSafetiesSurfBVH<<<initBlocks, initThreads>>>(points, origStates, bvhSafeties, config);
  VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
  time_safety_surf_bvh = timer.Stop();

  if (!config.only_surf) {
    if (config.validate_results) {
      ValidateSafety<<<initBlocks, initThreads>>>(safeties, refSafeties, num_better_safety_d, num_worse_safety_d,
                                                  config);
      VECGEOM_DEVICE_API_CALL(
          Memcpy(&num_better_safety, num_better_safety_d, sizeof(int), VECGEOM_DEVICE_API_SYMBOL(MemcpyDeviceToHost)));
      VECGEOM_DEVICE_API_CALL(
          Memcpy(&num_worse_safety, num_worse_safety_d, sizeof(int), VECGEOM_DEVICE_API_SYMBOL(MemcpyDeviceToHost)));
      VECGEOM_DEVICE_API_CALL(DeviceSynchronize());

      if (num_better_safety > 0)
        std::cout << "Device:    number of better safety values: " << num_better_safety << "\n";
      if (num_worse_safety > 0) std::cout << "Device:    number of worse safety values: " << num_worse_safety << "\n";

      if (config.test_bvh) {
        // Init counters
        num_better_safety = num_worse_safety = 0;
        VECGEOM_DEVICE_API_CALL(Memcpy(num_better_safety_d, &num_better_safety, sizeof(int),
                                       VECGEOM_DEVICE_API_SYMBOL(MemcpyHostToDevice)));
        VECGEOM_DEVICE_API_CALL(
            Memcpy(num_worse_safety_d, &num_worse_safety, sizeof(int), VECGEOM_DEVICE_API_SYMBOL(MemcpyHostToDevice)));

        ValidateSafety<<<initBlocks, initThreads>>>(bvhSafeties, refSafeties, num_better_safety_d, num_worse_safety_d,
                                                    config);
        VECGEOM_DEVICE_API_CALL(Memcpy(&num_better_safety, num_better_safety_d, sizeof(int),
                                       VECGEOM_DEVICE_API_SYMBOL(MemcpyDeviceToHost)));
        VECGEOM_DEVICE_API_CALL(
            Memcpy(&num_worse_safety, num_worse_safety_d, sizeof(int), VECGEOM_DEVICE_API_SYMBOL(MemcpyDeviceToHost)));
        VECGEOM_DEVICE_API_CALL(DeviceSynchronize());

        if (num_better_safety > 0)
          std::cout << "Device:    BVH number of better safety values: " << num_better_safety << "\n";
        if (num_worse_safety > 0)
          std::cout << "Device:    BVH number of worse safety values: " << num_worse_safety << "\n";
      }
    }
  }

  if (!config.debug) {
    if (config.only_surf)
      std::cout << "Device: safety_surf: " << time_safety_surf;
    else
      std::cout << "Device: safety_solids: " << time_safety_solids << " safety_solids_BVH: " << time_safety_solids_bvh
                << "  safety_surf: " << time_safety_surf;
    if (config.test_bvh) {
      std::cout << "  safety_surf_bvh: " << time_safety_surf_bvh << "\n";
    } else {
      std::cout << "\n";
    }
  }

  if (!config.only_surf) {
    // Traversal for solids model (reference)
    timer.Start();
    PropagateRaysSolid<LoopNavigator>
        <<<initBlocks, initThreads>>>(points, dirs, origStates, refLength_over_crossings, config);
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
    time_traverse_solids = timer.Stop();

    // Traversal for solids model with BVH
    timer.Start();
    PropagateRaysSolid<BVHNavigator>
        <<<initBlocks, initThreads>>>(points, dirs, origStates, length_over_crossings, config);
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
    time_traverse_solids_bvh = timer.Stop();
  }

  // Traversal for the surface model
  timer.Start();
  PropagateRaysSurf<<<initBlocks, initThreads>>>(points, dirs, origStates, length_over_crossings, surfdata_D, false,
                                                 config);
  VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
  time_traverse_surf = timer.Stop();

  // Traversal for the surface model with BVH
  timer.Start();
  if (config.test_bvh)
    PropagateRaysSurfBVH<<<initBlocks, initThreads>>>(points, dirs, origStates, length_over_crossings_bvh, surfdata_D,
                                                      false, config);
  VECGEOM_DEVICE_API_CALL(DeviceSynchronize());
  time_traverse_surf_bvh = timer.Stop();

  // Traversal for the surface model with BVH in single step mode
  NavigationState *BVH_single_states;
  VECGEOM_DEVICE_API_CALL(Malloc(&BVH_single_states, config.nrays * sizeof(NavigationState)));
  VECGEOM_DEVICE_API_CALL(Memcpy(BVH_single_states, origStates, config.nrays * sizeof(NavigationState),
                                 VECGEOM_DEVICE_API_SYMBOL(MemcpyDeviceToDevice)));

  // Traversal for the surface model with BVH in a single step mode, each step is done in a separate kernel launch.
  timer.Start();
  if (config.bvh_single_step)
    PropagateRaysSurfBVH(initThreads, points_NC, dirs, BVH_single_states, length_over_crossings_bvh, false, config);
  auto time_traverse_surf_bvh_single = timer.Stop();

  // Traversal for the surface model with BVH in a single step mode, each step is done in a separate kernel launch.
  // The transport kernel is split, one for identifying the framed surface, one for the relocation on the common surface
  timer.Start();
  if (config.bvh_split_step) {
    // reset points and in_states, as they were changed by the single step kernel above
    VECGEOM_DEVICE_API_CALL(Memcpy(BVH_single_states, origStates, config.nrays * sizeof(NavigationState),
                                   VECGEOM_DEVICE_API_SYMBOL(MemcpyDeviceToDevice)));
    VECGEOM_DEVICE_API_CALL(
        Memcpy(points_NC, pointsh, config.nrays * sizeof(Vec3D), VECGEOM_DEVICE_API_SYMBOL(MemcpyHostToDevice)));
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());

    PropagateRaysSurfBVHSplit(initThreads, points_NC, dirs, BVH_single_states, length_over_crossings_bvh, false,
                              config);
  }
  auto time_traverse_surf_bvh_single_split = timer.Stop();

  if (!config.only_surf && config.validate_results) {
    auto world_dev = vecgeom::cxx::CudaManager::Instance().world_gpu();
    ValidateTraversal<<<initBlocks, initThreads>>>(points, dirs, origStates, length_over_crossings,
                                                   refLength_over_crossings, num_errors_dist_d, world_dev, surfdata_D,
                                                   /*use_bvh=*/false, config);
    VECGEOM_DEVICE_API_CALL(
        Memcpy(&num_errors_dist, num_errors_dist_d, sizeof(int), VECGEOM_DEVICE_API_SYMBOL(MemcpyDeviceToHost)));
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());

    if (config.test_bvh)
      ValidateTraversal<<<initBlocks, initThreads>>>(points, dirs, origStates, length_over_crossings_bvh,
                                                     refLength_over_crossings, num_errors_dist_bvh_d, world_dev,
                                                     surfdata_D, /*use_bvh=*/true, config);
    VECGEOM_DEVICE_API_CALL(Memcpy(&num_errors_dist_bvh, num_errors_dist_bvh_d, sizeof(int),
                                   VECGEOM_DEVICE_API_SYMBOL(MemcpyDeviceToHost)));
    VECGEOM_DEVICE_API_CALL(DeviceSynchronize());

    if (num_errors_dist > 0 || num_errors_dist_bvh > 0)
      std::cout << "Device: traversal errors looper: " << num_errors_dist
                << " traversal errors BVH: " << num_errors_dist_bvh << "\n";
  }

  if (!config.debug) {
    if (config.only_surf)
      std::cout << "Device: traverse_surf: " << time_traverse_surf;
    else
      std::cout << "Device: traverse_solids: " << time_traverse_solids
                << " traverse_solids_bvh: " << time_traverse_solids_bvh << "  traverse_surf: " << time_traverse_surf;
    if (config.test_bvh) std::cout << "  traverse_surf BVH: " << time_traverse_surf_bvh;
    if (config.bvh_single_step) std::cout << "  traverse_surf BVH single steps: " << time_traverse_surf_bvh_single;
    if (config.bvh_split_step)
      std::cout << "  traverse_surf BVH single split steps: " << time_traverse_surf_bvh_single_split;
    std::cout << "\n";
  }

  BrepCudaManager::Instance().Cleanup();
  delete[] pointsh;
  delete[] dirsh;
  if (!config.only_surf) {
    num_errors += num_errors_dist + num_errors_dist_bvh;
    if (num_errors > 0) printf("Device: num_erros = %d / %d\n", num_errors, config.nrays);
    VECGEOM_DEVICE_API_CALL(Free(num_errors_d));
    VECGEOM_DEVICE_API_CALL(Free(num_better_safety_d));
    VECGEOM_DEVICE_API_CALL(Free(num_worse_safety_d));
    VECGEOM_DEVICE_API_CALL(Free(num_errors_dist_d));
    VECGEOM_DEVICE_API_CALL(Free(num_errors_dist_bvh_d));
  }
  VECGEOM_DEVICE_API_CALL(Free(points));
  VECGEOM_DEVICE_API_CALL(Free(dirs));
  VECGEOM_DEVICE_API_CALL(Free(points_NC));
  VECGEOM_DEVICE_API_CALL(Free(origStates));
  VECGEOM_DEVICE_API_CALL(Free(outputStates));
  VECGEOM_DEVICE_API_CALL(Free(outputStatesBVH));
  if (config.validate_results) // These arrays will be nullptr if we are not doing validation
  {
    VECGEOM_DEVICE_API_CALL(Free(refSafeties));
    VECGEOM_DEVICE_API_CALL(Free(safeties));
    VECGEOM_DEVICE_API_CALL(Free(refLength_over_crossings));
    VECGEOM_DEVICE_API_CALL(Free(length_over_crossings));
    VECGEOM_DEVICE_API_CALL(Free(length_over_crossings_bvh));
  }
  if (config.only_surf) return 0;
  return num_errors;
}
